from langchain_core.messages import HumanMessage
from langgraph.graph import MessageGraph, END


def entry(input:list[HumanMessage]):
    return input
def one_task(input:list[HumanMessage]):
    print("working with one task")
    return input
def two_task(input:list[HumanMessage]):
    print("working with two task")
    return input
def router(input:list[HumanMessage]):
    if "two task" in input[0].content:
        return "two_task"
    else:
        return "one_task"


workflow = MessageGraph()

workflow.add_node("entry", entry)
workflow.add_node("two_task", two_task)
workflow.add_node("one_task", one_task)

# 条件边
workflow.add_conditional_edges(
    "entry", router, {"two_task": "two_task", "one_task": "one_task"}
)
# 普通边
workflow.add_edge("two_task", END)
workflow.add_edge("one_task", END)

workflow.set_entry_point("entry")

app = workflow.compile()
res = app.invoke("I want to two task")
print(res)
